import gradio as gr
import numpy as np
import pickle
import cv2
import logging
import os
import subprocess
import sys

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

class CatDogClassifier:
    def __init__(self):
        """初始化猫狗分类器"""
        self.model = None
        self.model_loaded = False
        self.load_model()
        
    def load_model(self):
        """加载训练好的模型"""
        try:
            if os.path.exists("best_ensemble_model.pkl"):
                with open("best_ensemble_model.pkl", 'rb') as f:
                    self.model = pickle.load(f)
                self.model_loaded = True
                logging.info("模型加载成功")
            else:
                logging.warning("模型文件不存在，请先训练模型")
                self.model_loaded = False
        except Exception as e:
            logging.error(f"加载模型失败: {e}")
            self.model_loaded = False
    
    def train_model(self):
        """训练模型"""
        try:
            logging.info("开始训练模型...")
            # 调用训练脚本
            result = subprocess.run([sys.executable, "train_and_deploy.py"], 
                                  capture_output=True, text=True)
            
            if result.returncode == 0:
                logging.info("模型训练成功")
                # 重新加载模型
                self.load_model()
                return "模型训练成功！现在可以上传图片进行分类了。"
            else:
                error_msg = f"训练失败: {result.stderr}"
                logging.error(error_msg)
                return error_msg
                
        except Exception as e:
            error_msg = f"训练过程中出错: {e}"
            logging.error(error_msg)
            return error_msg
    
    def preprocess_image(self, input_image):
        """
        预处理上传的图片
        
        参数:
            input_image: 输入的图片数据
            
        返回:
            numpy.array: 预处理后的特征向量
        """
        try:
            # 将图片转换为numpy数组
            if isinstance(input_image, str):
                # 如果是文件路径，读取图片
                img_array = cv2.imread(input_image)
            else:
                # 如果是numpy数组
                img_array = input_image
            
            if img_array is None:
                raise ValueError("无法读取图片")
            
            # 使用flat特征提取方法（与训练时一致）
            if len(img_array.shape) == 3:
                img_gray = cv2.cvtColor(img_array, cv2.COLOR_BGR2GRAY)
            else:
                img_gray = img_array
                
            img_resized = cv2.resize(img_gray, (32, 32))
            img_flatten = img_resized.flatten()
            img_normalized = img_flatten.astype('float32') / 255.0
            return img_normalized.reshape(1, -1)
                
        except Exception as e:
            logging.error(f"图片预处理失败: {e}")
            raise e
    
    def predict(self, input_image):
        """
        预测图片是猫还是狗
        
        参数:
            input_image: 输入的图片
            
        返回:
            dict: 预测结果和置信度
        """
        if not self.model_loaded:
            return {
                'prediction': '模型未加载',
                'confidence': 0.0,
                'probabilities': {'猫': 0.0, '狗': 0.0},
                'error': '请先训练模型'
            }
        
        try:
            # 预处理图片
            processed_image = self.preprocess_image(input_image)
            
            # 进行预测
            if hasattr(self.model, 'predict_proba'):
                prediction = self.model.predict(processed_image)[0]
                probability = self.model.predict_proba(processed_image)[0]
            else:
                # 对于没有predict_proba的模型（如硬投票）
                prediction = self.model.predict(processed_image)[0]
                probability = [0.5, 0.5]  # 默认概率
                logging.warning("当前模型不支持概率预测")
            
            # 获取预测标签和置信度
            class_names = ['猫', '狗']
            predicted_class = class_names[prediction]
            confidence = probability[prediction] if len(probability) > prediction else 0.5
            
            # 返回结果
            result = {
                'prediction': predicted_class,
                'confidence': float(confidence),
                'probabilities': {
                    '猫': float(probability[0]) if len(probability) > 0 else 0.5,
                    '狗': float(probability[1]) if len(probability) > 1 else 0.5
                }
            }
            
            logging.info(f"预测结果: {result}")
            return result
            
        except Exception as e:
            logging.error(f"预测失败: {e}")
            return {
                'prediction': '错误',
                'confidence': 0.0,
                'probabilities': {'猫': 0.0, '狗': 0.0},
                'error': str(e)
            }

# 创建分类器实例
classifier = CatDogClassifier()

def classify_image(input_image):
    """
    Gradio接口函数，用于分类上传的图片
    """
    if not classifier.model_loaded:
        return "请先训练模型", 0, 0
    
    try:
        # 进行预测
        result = classifier.predict(input_image)
        
        # 如果有错误，返回错误信息
        if 'error' in result:
            return f"预测错误: {result['error']}", 0, 0
        
        # 格式化输出
        prediction_text = f"预测结果: {result['prediction']} (置信度: {result['confidence']:.2%})"
        
        # 返回结果
        return prediction_text, result['probabilities']['猫'], result['probabilities']['狗']
        
    except Exception as e:
        error_msg = f"处理图片时发生错误: {str(e)}"
        logging.error(error_msg)
        return error_msg, 0, 0

def train_model_interface():
    """训练模型的接口函数"""
    result = classifier.train_model()
    return result

def create_interface():
    """创建Gradio界面"""
    
    # 界面描述
    description = """
    # 🐱🐶 猫狗分类器
    
    上传一张猫或狗的图片，AI模型会自动识别图片中是猫还是狗！
    
    **使用说明:**
    1. 首先点击"训练模型"按钮（只需要训练一次）
    2. 上传猫狗图片或使用示例图片
    3. 查看分类结果和概率
    
    **注意:** 请确保数据目录 `./data/train/` 中包含猫狗图片数据。
    """
    
    # 创建界面
    with gr.Blocks(theme=gr.themes.Soft(), title="猫狗分类器") as interface:
        gr.Markdown(description)
        
        with gr.Row():
            with gr.Column():
                # 训练模型按钮
                train_btn = gr.Button("训练模型", variant="primary")
                train_output = gr.Textbox(label="训练状态", interactive=False)
                
                # 图片输入组件
                image_input = gr.Image(
                    label="上传图片",
                    type="numpy",
                    source="upload",
                    interactive=True
                )
                
            with gr.Column():
                # 结果显示
                result_output = gr.Textbox(
                    label="预测结果",
                    placeholder="这里将显示预测结果...",
                    interactive=False
                )
                
                # 概率显示
                with gr.Row():
                    cat_prob = gr.Number(
                        label="猫的概率",
                        minimum=0,
                        maximum=1,
                        precision=4
                    )
                    
                    dog_prob = gr.Number(
                        label="狗的概率", 
                        minimum=0,
                        maximum=1,
                        precision=4
                    )
        
        # 示例图片
        gr.Examples(
            examples=[
                ["examples/cat_example.jpg"] if os.path.exists("examples/cat_example.jpg") else ["https://images.unsplash.com/photo-1514888286974-6c03e2ca1dba?w=200"],
                ["examples/dog_example.jpg"] if os.path.exists("examples/dog_example.jpg") else ["https://images.unsplash.com/photo-1561037404-61cd46aa615b?w=200"]
            ],
            inputs=image_input,
            label="示例图片 (点击试试)"
        )
        
        # 底部信息
        gr.Markdown("""
        ---
        **技术说明:**
        - 使用集成学习模型进行预测
        - 特征提取方法: Flat (灰度+32x32分辨率)
        - 支持格式: JPG, PNG, JPEG等常见图片格式
        """)
        
        # 绑定事件
        train_btn.click(
            fn=train_model_interface,
            outputs=[train_output]
        )
        
        # 当图片输入变化时触发预测
        image_input.change(
            fn=classify_image,
            inputs=[image_input],
            outputs=[result_output, cat_prob, dog_prob]
        )
    
    return interface

if __name__ == "__main__":
    # 创建界面
    interface = create_interface()
    
    # 启动服务
    logging.info("启动Gradio服务...")
    interface.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        inbrowser=True
    )